{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "kernelspec": {
      "display_name": "Python 3",
      "language": "python",
      "name": "python3"
    },
    "language_info": {
      "codemirror_mode": {
        "name": "ipython",
        "version": 3
      },
      "file_extension": ".py",
      "mimetype": "text/x-python",
      "name": "python",
      "nbconvert_exporter": "python",
      "pygments_lexer": "ipython3",
      "version": "3.8.3"
    },
    "colab": {
      "name": "HW05_ZH.ipynb",
      "provenance": []
    }
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "s4uNq4OUsqkS"
      },
      "source": [
        "# **Homework 5 - Sequence-to-sequence**\n",
        "\n",
        "若有任何問題，歡迎來信至助教信箱 ntu-ml-2021spring-ta@googlegroups.com\n",
        "\n",
        "- Link to [tutorial]() (TBA)."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "a4b3pTBNsqke"
      },
      "source": [
        "# Sequence-to-Sequence 介紹\n",
        "- 大多數常見的 seq2seq model 為 encoder-decoder model，主要由兩個部分組成，分別是 encoder 和 decoder，而這兩個部可以使用 recurrent neural network (RNN)或 transformer 來實作，主要是用來解決輸入和輸出的長度不一樣的情況\n",
        "- **Encoder** 是將一連串的輸入，如文字、影片、聲音訊號等，編碼為單個向量，這單個向量可以想像為是整個輸入的抽象表示，包含了整個輸入的資訊\n",
        "- **Decoder** 是將 encoder 輸出的單個向量逐步解碼，一次輸出一個結果，直到將最後目標輸出被產生出來為止，每次輸出會影響下一次的輸出，一般會在開頭加入 \"< BOS >\" 來表示開始解碼，會在結尾輸出 \"< EOS >\" 來表示輸出結束\n",
        "\n",
        "\n",
        "![seq2seq](https://i.imgur.com/0zeDyuI.png)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "xBtwjLeZsqkg"
      },
      "source": [
        "# 作業介紹\n",
        "- 英文翻譯中文\n",
        "  - 輸入： 一句英文 （e.g.\t\ttom is a student .） \n",
        "  - 輸出： 中文翻譯 （e.g. \t\t湯姆 是 個 學生 。）\n",
        "\n",
        "- TODO\n",
        "  - 訓練一個 RNN 模型達到 Seq2seq 翻譯\n",
        "  - 訓練一個 Transformer 大幅提升效能\n",
        "  - 實作 Back-translation 大幅提升效能"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ujbwECV9sqkh"
      },
      "source": [
        "# 下載和引入需要的函式庫"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "z0rSekkNsqkh"
      },
      "source": [
        "!pip install 'torch>=1.6.0' editdistance matplotlib sacrebleu sacremoses sentencepiece tqdm wandb\n",
        "!pip install --upgrade jupyter ipywidgets"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "QKgTjk7Psqki"
      },
      "source": [
        "!git clone https://github.com/pytorch/fairseq.git\n",
        "!cd fairseq && git checkout 9a1c497\n",
        "!pip install --upgrade ./fairseq/"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "q2hPnvIpsqkj"
      },
      "source": [
        "import sys\n",
        "import pdb\n",
        "import pprint\n",
        "import logging\n",
        "import os\n",
        "import random\n",
        "\n",
        "import torch\n",
        "import torch.nn as nn\n",
        "import torch.nn.functional as F\n",
        "from torch.utils import data\n",
        "import numpy as np\n",
        "import tqdm.auto as tqdm\n",
        "from pathlib import Path\n",
        "from argparse import Namespace\n",
        "from fairseq import utils\n",
        "\n",
        "import matplotlib.pyplot as plt"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "0SokhAvEsqkj"
      },
      "source": [
        "# 設定種子"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "JtsGxPFZsqkk"
      },
      "source": [
        "seed = 73\n",
        "random.seed(seed)\n",
        "torch.manual_seed(seed)\n",
        "if torch.cuda.is_available():\n",
        "    torch.cuda.manual_seed(seed)\n",
        "    torch.cuda.manual_seed_all(seed)  \n",
        "np.random.seed(seed)  \n",
        "torch.backends.cudnn.benchmark = False\n",
        "torch.backends.cudnn.deterministic = True"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Lpx0mMKvsqkk"
      },
      "source": [
        "# 資料集介紹\n",
        "\n",
        "## 英轉繁雙語資料\n",
        "* [TED2020](#reimers-2020-multilingual-sentence-bert)\n",
        "    - 原始資料量: 398,066句    \n",
        "    - 處理後資料: 393,980句\n",
        "    \n",
        "\n",
        "## 測試資料\n",
        "- 資料量: 4,000句\n",
        "- **中文部分不公開，提供的檔案為假翻譯，全部都是句點。**"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "q6ICndcwsqkl"
      },
      "source": [
        "# 資料下載"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Rz8ejiDysqkl"
      },
      "source": [
        "### 安裝megatools (optional)"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "speRSkeosqkl"
      },
      "source": [
        "#!apt-get install megatools"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "oBoh6yBksqkm"
      },
      "source": [
        "## 下載檔案並解壓縮"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "Qdb_grqAsqkm"
      },
      "source": [
        "data_dir = './DATA/rawdata'\n",
        "dataset_name = 'ted2020'\n",
        "urls = (\n",
        "    '\"https://onedrive.live.com/download?cid=3E549F3B24B238B4&resid=3E549F3B24B238B4%214989&authkey=AGgQ-DaR8eFSl1A\"', \n",
        "    '\"https://onedrive.live.com/download?cid=3E549F3B24B238B4&resid=3E549F3B24B238B4%214987&authkey=AA4qP_azsicwZZM\"',\n",
        "# # If the above links die, use the following instead. \n",
        "#     \"https://www.csie.ntu.edu.tw/~r09922057/ML2021-hw5/ted2020.tgz\",\n",
        "#     \"https://www.csie.ntu.edu.tw/~r09922057/ML2021-hw5/test.tgz\",\n",
        "# # If the above links die, use the following instead. \n",
        "#     \"https://mega.nz/#!vEcTCISJ!3Rw0eHTZWPpdHBTbQEqBDikDEdFPr7fI8WxaXK9yZ9U\",\n",
        "#     \"https://mega.nz/#!zNcnGIoJ!oPJX9AvVVs11jc0SaK6vxP_lFUNTkEcK2WbxJpvjU5Y\",\n",
        ")\n",
        "file_names = (\n",
        "    'ted2020.tgz', # train & dev\n",
        "    'test.tgz', # test\n",
        ")\n",
        "prefix = Path(data_dir).absolute() / dataset_name\n",
        "\n",
        "prefix.mkdir(parents=True, exist_ok=True)\n",
        "for u, f in zip(urls, file_names):\n",
        "    path = prefix/f\n",
        "    if not path.exists():\n",
        "        if 'mega' in u:\n",
        "            !megadl {u} --path {path}\n",
        "        else:\n",
        "            !wget {u} -O {path}\n",
        "    if path.suffix == \".tgz\":\n",
        "        !tar -xvf {path} -C {prefix}\n",
        "    elif path.suffix == \".zip\":\n",
        "        !unzip -o {path} -d {prefix}\n",
        "!mv {prefix/'raw.en'} {prefix/'train_dev.raw.en'}\n",
        "!mv {prefix/'raw.zh'} {prefix/'train_dev.raw.zh'}\n",
        "!mv {prefix/'test.en'} {prefix/'test.raw.en'}\n",
        "!mv {prefix/'test.zh'} {prefix/'test.raw.zh'}"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "djXV5NPFsqkn"
      },
      "source": [
        "## 設定語言"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "ld-Bphtrsqkn"
      },
      "source": [
        "src_lang = 'en'\n",
        "tgt_lang = 'zh'\n",
        "\n",
        "data_prefix = f'{prefix}/train_dev.raw'\n",
        "test_prefix = f'{prefix}/test.raw'"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "4gejawdXsqko"
      },
      "source": [
        "!head {data_prefix+'.'+src_lang} -n 5\n",
        "!head {data_prefix+'.'+tgt_lang} -n 5"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "BvH3snDqsqko"
      },
      "source": [
        "## 檔案前處理"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "2GZQLH9ksqko"
      },
      "source": [
        "import re\n",
        "\n",
        "def strQ2B(ustring):\n",
        "    \"\"\"把字串全形轉半形\"\"\"\n",
        "    # 參考來源:https://ithelp.ithome.com.tw/articles/10233122\n",
        "    ss = []\n",
        "    for s in ustring:\n",
        "        rstring = \"\"\n",
        "        for uchar in s:\n",
        "            inside_code = ord(uchar)\n",
        "            if inside_code == 12288:  # 全形空格直接轉換\n",
        "                inside_code = 32\n",
        "            elif (inside_code >= 65281 and inside_code <= 65374):  # 全形字元（除空格）根據關係轉化\n",
        "                inside_code -= 65248\n",
        "            rstring += chr(inside_code)\n",
        "        ss.append(rstring)\n",
        "    return ''.join(ss)\n",
        "                \n",
        "def clean_s(s, lang):\n",
        "    if lang == 'en':\n",
        "        s = re.sub(r\"\\([^()]*\\)\", \"\", s) # remove ([text])\n",
        "        s = s.replace('-', '') # remove '-'\n",
        "        s = re.sub('([.,;!?()\\\"])', r' \\1 ', s) # keep punctuation\n",
        "    elif lang == 'zh':\n",
        "        s = strQ2B(s) # Q2B\n",
        "        s = re.sub(r\"\\([^()]*\\)\", \"\", s) # remove ([text])\n",
        "        s = s.replace(' ', '')\n",
        "        s = s.replace('—', '')\n",
        "        s = s.replace('“', '\"')\n",
        "        s = s.replace('”', '\"')\n",
        "        s = s.replace('_', '')\n",
        "        s = re.sub('([。,;!?()\\\"~「」])', r' \\1 ', s) # keep punctuation\n",
        "    s = ' '.join(s.strip().split())\n",
        "    return s\n",
        "\n",
        "def len_s(s, lang):\n",
        "    if lang == 'zh':\n",
        "        return len(s)\n",
        "    return len(s.split())\n",
        "\n",
        "def clean_corpus(prefix, l1, l2, ratio=9, max_len=1000, min_len=1):\n",
        "    if Path(f'{prefix}.clean.{l1}').exists() and Path(f'{prefix}.clean.{l2}').exists():\n",
        "        print(f'{prefix}.clean.{l1} & {l2} exists. skipping clean.')\n",
        "        return\n",
        "    with open(f'{prefix}.{l1}', 'r') as l1_in_f:\n",
        "        with open(f'{prefix}.{l2}', 'r') as l2_in_f:\n",
        "            with open(f'{prefix}.clean.{l1}', 'w') as l1_out_f:\n",
        "                with open(f'{prefix}.clean.{l2}', 'w') as l2_out_f:\n",
        "                    for s1 in l1_in_f:\n",
        "                        s1 = s1.strip()\n",
        "                        s2 = l2_in_f.readline().strip()\n",
        "                        s1 = clean_s(s1, l1)\n",
        "                        s2 = clean_s(s2, l2)\n",
        "                        s1_len = len_s(s1, l1)\n",
        "                        s2_len = len_s(s2, l2)\n",
        "                        if min_len > 0: # remove short sentence\n",
        "                            if s1_len < min_len or s2_len < min_len:\n",
        "                                continue\n",
        "                        if max_len > 0: # remove long sentence\n",
        "                            if s1_len > max_len or s2_len > max_len:\n",
        "                                continue\n",
        "                        if ratio > 0: # remove by ratio of length\n",
        "                            if s1_len/s2_len > ratio or s2_len/s1_len > ratio:\n",
        "                                continue\n",
        "                        print(s1, file=l1_out_f)\n",
        "                        print(s2, file=l2_out_f)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "1i1JadfRsqks"
      },
      "source": [
        "clean_corpus(data_prefix, src_lang, tgt_lang)\n",
        "clean_corpus(test_prefix, src_lang, tgt_lang, ratio=-1, min_len=-1, max_len=-1)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "1wrLAc41sqkt"
      },
      "source": [
        "!head {data_prefix+'.clean.'+src_lang} -n 5\n",
        "!head {data_prefix+'.clean.'+tgt_lang} -n 5"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "GNXEL5jTsqku"
      },
      "source": [
        "## 切出 train/valid set"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "wi5OiMDDsqkv"
      },
      "source": [
        "valid_ratio = 0.01 # 3000~4000句就夠了\n",
        "train_ratio = 1 - valid_ratio"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "bUx2nQsJsqkv"
      },
      "source": [
        "if (prefix/f'train.clean.{src_lang}').exists() \\\n",
        "and (prefix/f'train.clean.{tgt_lang}').exists() \\\n",
        "and (prefix/f'valid.clean.{src_lang}').exists() \\\n",
        "and (prefix/f'valid.clean.{tgt_lang}').exists():\n",
        "    print(f'train/valid splits exists. skipping split.')\n",
        "else:\n",
        "    line_num = sum(1 for line in open(f'{data_prefix}.clean.{src_lang}'))\n",
        "    labels = list(range(line_num))\n",
        "    random.shuffle(labels)\n",
        "    for lang in [src_lang, tgt_lang]:\n",
        "        train_f = open(os.path.join(data_dir, dataset_name, f'train.clean.{lang}'), 'w')\n",
        "        valid_f = open(os.path.join(data_dir, dataset_name, f'valid.clean.{lang}'), 'w')\n",
        "        count = 0\n",
        "        for line in open(f'{data_prefix}.clean.{lang}', 'r'):\n",
        "            if labels[count]/line_num < train_ratio:\n",
        "                train_f.write(line)\n",
        "            else:\n",
        "                valid_f.write(line)\n",
        "            count += 1\n",
        "        train_f.close()\n",
        "        valid_f.close()"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Bpkexvlgsqkv"
      },
      "source": [
        "## Subword Units \n",
        "翻譯存在的一大問題是未登錄詞(out of vocabulary)，可以使用 subword units 作為斷詞單位來解決。\n",
        "- 使用 [sentencepiece](#kudo-richardson-2018-sentencepiece) 套件\n",
        "- 用 unigram 或 byte-pair encoding (BPE)"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "Fd700Px8sqkw"
      },
      "source": [
        "import sentencepiece as spm\n",
        "vocab_size = 8000\n",
        "if (prefix/f'spm{vocab_size}.model').exists():\n",
        "    print(f'{prefix}/spm{vocab_size}.model exists. skipping spm_train.')\n",
        "else:\n",
        "    spm.SentencePieceTrainer.train(\n",
        "        input=','.join([f'{prefix}/train.clean.{src_lang}',\n",
        "                        f'{prefix}/valid.clean.{src_lang}',\n",
        "                        f'{prefix}/train.clean.{tgt_lang}',\n",
        "                        f'{prefix}/valid.clean.{tgt_lang}']),\n",
        "        model_prefix=prefix/f'spm{vocab_size}',\n",
        "        vocab_size=vocab_size,\n",
        "        character_coverage=1,\n",
        "        model_type='unigram', # 'bpe' 也可\n",
        "        input_sentence_size=1e6,\n",
        "        shuffle_input_sentence=True,\n",
        "        normalization_rule_name='nmt_nfkc_cf',\n",
        "    )"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "knA04up1sqkx"
      },
      "source": [
        "spm_model = spm.SentencePieceProcessor(model_file=str(prefix/f'spm{vocab_size}.model'))\n",
        "in_tag = {\n",
        "    'train': 'train.clean',\n",
        "    'valid': 'valid.clean',\n",
        "    'test': 'test.raw.clean',\n",
        "}\n",
        "for split in ['train', 'valid', 'test']:\n",
        "    for lang in [src_lang, tgt_lang]:\n",
        "        out_path = prefix/f'{split}.{lang}'\n",
        "        if out_path.exists():\n",
        "            print(f\"{out_path} exists. skipping spm_encode.\")\n",
        "        else:\n",
        "            with open(prefix/f'{split}.{lang}', 'w') as out_f:\n",
        "                with open(prefix/f'{in_tag[split]}.{lang}', 'r') as in_f:\n",
        "                    for line in in_f:\n",
        "                        line = line.strip()\n",
        "                        tok = spm_model.encode(line, out_type=str)\n",
        "                        print(' '.join(tok), file=out_f)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "D9YEuLYZsqkx"
      },
      "source": [
        "!head {data_dir+'/'+dataset_name+'/train.'+src_lang} -n 5\n",
        "!head {data_dir+'/'+dataset_name+'/train.'+tgt_lang} -n 5"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "cY-7ym0Nsqky"
      },
      "source": [
        "## 用 fairseq 將資料轉為 binary"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "H4uvbYrusqky"
      },
      "source": [
        "binpath = Path('./DATA/data-bin', dataset_name)\n",
        "if binpath.exists():\n",
        "    print(binpath, \"exists, will not overwrite!\")\n",
        "else:\n",
        "    !python -m fairseq_cli.preprocess \\\n",
        "        --source-lang {src_lang}\\\n",
        "        --target-lang {tgt_lang}\\\n",
        "        --trainpref {prefix/'train'}\\\n",
        "        --validpref {prefix/'valid'}\\\n",
        "        --testpref {prefix/'test'}\\\n",
        "        --destdir {binpath}\\\n",
        "        --joined-dictionary\\\n",
        "        --workers 2"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "zln6Zku8sqky"
      },
      "source": [
        "# 實驗的參數設定表"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "yzEG28ypsqkz"
      },
      "source": [
        "config = Namespace(\n",
        "    datadir = \"./DATA/data-bin/ted2020\",\n",
        "    savedir = \"./checkpoints/rnn\",\n",
        "    source_lang = \"en\",\n",
        "    target_lang = \"zh\",\n",
        "    \n",
        "    # cpu threads when fetching & processing data.\n",
        "    num_workers=2,  \n",
        "    # batch size in terms of tokens. gradient accumulation increases the effective batchsize.\n",
        "    max_tokens=8192,\n",
        "    accum_steps=2,\n",
        "    \n",
        "    # the lr s calculated from Noam lr scheduler. you can tune the maximum lr by this factor.\n",
        "    lr_factor=2.,\n",
        "    lr_warmup=4000,\n",
        "    \n",
        "    # clipping gradient norm helps alleviate gradient exploding\n",
        "    clip_norm=1.0,\n",
        "    \n",
        "    # maximum epochs for training\n",
        "    max_epoch=30,\n",
        "    start_epoch=1,\n",
        "    \n",
        "    # beam size for beam search\n",
        "    beam=5, \n",
        "    # generate sequences of maximum length ax + b, where x is the source length\n",
        "    max_len_a=1.2, \n",
        "    max_len_b=10,\n",
        "    # when decoding, post process sentence by removing sentencepiece symbols.\n",
        "    post_process = \"sentencepiece\",\n",
        "    \n",
        "    # checkpoints\n",
        "    keep_last_epochs=5,\n",
        "    resume=None, # if resume from checkpoint name (under config.savedir)\n",
        "    \n",
        "    # logging\n",
        "    use_wandb=False,\n",
        ")"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "OVJ6D_T3sqkz"
      },
      "source": [
        "# Logging\n",
        "- logging 套件紀錄一般訊息\n",
        "- wandb 紀錄續練過程 loss, bleu, model weight 等等"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "scrolled": true,
        "id": "jP-CfPkosqkz"
      },
      "source": [
        "logging.basicConfig(\n",
        "    format=\"%(asctime)s | %(levelname)s | %(name)s | %(message)s\",\n",
        "    datefmt=\"%Y-%m-%d %H:%M:%S\",\n",
        "    level=\"INFO\", # \"DEBUG\" \"WARNING\" \"ERROR\"\n",
        "    stream=sys.stdout,\n",
        ")\n",
        "proj = \"hw5.seq2seq\"\n",
        "logger = logging.getLogger(proj)\n",
        "if config.use_wandb:\n",
        "    import wandb\n",
        "    wandb.init(project=proj, name=Path(config.savedir).stem, config=config)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "WcWmOCMPsqk0"
      },
      "source": [
        "# CUDA環境"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "6L2wRcbYsqk0"
      },
      "source": [
        "cuda_env = utils.CudaEnvironment()\n",
        "utils.CudaEnvironment.pretty_print_cuda_env_list([cuda_env])\n",
        "device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "zUth0pCWsqk0"
      },
      "source": [
        "# 讀取資料集"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "MQacez3dsqk1"
      },
      "source": [
        "## 借用 fairseq 的 TranslationTask\n",
        "* 用來讀進上面 binarized 的檔案\n",
        "* 有現成的 data iterator (dataloader)\n",
        "* 字典 task.source_dictionary 和 task.target_dictionary 也很好用 \n",
        "* 有實做 beam search"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "A6a3rsGFsqk1"
      },
      "source": [
        "from fairseq.tasks.translation import TranslationConfig, TranslationTask\n",
        "\n",
        "## setup task\n",
        "task_cfg = TranslationConfig(\n",
        "    data=config.datadir,\n",
        "    source_lang=config.source_lang,\n",
        "    target_lang=config.target_lang,\n",
        "    train_subset=\"train\",\n",
        "    required_seq_len_multiple=8,\n",
        "    dataset_impl=\"mmap\",\n",
        "    upsample_primary=1,\n",
        ")\n",
        "task = TranslationTask.setup_task(task_cfg)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "SFxG6JpOsqk1"
      },
      "source": [
        "logger.info(\"loading data for epoch 1\")\n",
        "task.load_dataset(split=\"train\", epoch=1, combine=True) # combine if you have back-translation data.\n",
        "task.load_dataset(split=\"valid\", epoch=1)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "TyZ4mnSmsqk1"
      },
      "source": [
        "sample = task.dataset(\"valid\")[1]\n",
        "pprint.pprint(sample)\n",
        "pprint.pprint(\n",
        "    \"Source: \" + \\\n",
        "    task.source_dictionary.string(\n",
        "        sample['source'],\n",
        "        config.post_process,\n",
        "    )\n",
        ")\n",
        "pprint.pprint(\n",
        "    \"Target: \" + \\\n",
        "    task.target_dictionary.string(\n",
        "        sample['target'],\n",
        "        config.post_process,\n",
        "    )\n",
        ")"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "GDz-RErhsqk2"
      },
      "source": [
        "## Dataset Iterator"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "GZt8yQafsqk2"
      },
      "source": [
        "* 將每個 batch 控制在 N 個 token 讓 GPU 記憶體更有效被利用\n",
        "* 讓 training set 每個 epoch 有不同 shuffling\n",
        "* 濾掉長度太長的句子\n",
        "* 將每個 batch 內的句子 pad 成一樣長，好讓 GPU 平行運算\n",
        "* 加上 eos 並 shift 一格\n",
        "    - teacher forcing: 為了訓練模型根據prefix生成下個字，decoder的輸入會是輸出目標序列往右shift一格。\n",
        "    - 一般是會在輸入開頭加個bos token (如下圖)\n",
        "![seq2seq](https://i.imgur.com/0zeDyuI.png)\n",
        "    - fairseq 則是直接把 eos 挪到 beginning，訓練起來效果其實差不多。例如: \n",
        "    ```\n",
        "    # 輸出目標 (target) 和 Decoder輸入 (prev_output_tokens): \n",
        "                   eos = 2\n",
        "                target = 419,  711,  238,  888,  792,   60,  968,    8,    2\n",
        "    prev_output_tokens = 2,  419,  711,  238,  888,  792,   60,  968,    8\n",
        "    ```\n"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "Paco1lhXsqk2"
      },
      "source": [
        "def load_data_iterator(task, split, epoch=1, max_tokens=4000, num_workers=1, cached=True):\n",
        "    batch_iterator = task.get_batch_iterator(\n",
        "        dataset=task.dataset(split),\n",
        "        max_tokens=max_tokens,\n",
        "        max_sentences=None,\n",
        "        max_positions=utils.resolve_max_positions(\n",
        "            task.max_positions(),\n",
        "            max_tokens,\n",
        "        ),\n",
        "        ignore_invalid_inputs=True,\n",
        "        seed=seed,\n",
        "        num_workers=num_workers,\n",
        "        epoch=epoch,\n",
        "        disable_iterator_cache=not cached,\n",
        "        # Set this to False to speed up. However, if set to False, changing max_tokens beyond \n",
        "        # first call of this method has no effect. \n",
        "    )\n",
        "    return batch_iterator\n",
        "\n",
        "demo_epoch_obj = load_data_iterator(task, \"valid\", epoch=1, max_tokens=20, num_workers=1, cached=False)\n",
        "demo_iter = demo_epoch_obj.next_epoch_itr(shuffle=True)\n",
        "sample = next(demo_iter)\n",
        "sample"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "xuG0Youlsqk2"
      },
      "source": [
        "* 每個 batch 是一個字典，key 是字串，value 是 Tensor，內容說明如下\n",
        "```python\n",
        "batch = {\n",
        "    \"id\": id, # 每個 example 的 id\n",
        "    \"nsentences\": len(samples), # batch size 句子數\n",
        "    \"ntokens\": ntokens, # batch size 字數\n",
        "    \"net_input\": {\n",
        "        \"src_tokens\": src_tokens, # 來源語言的序列\n",
        "        \"src_lengths\": src_lengths, # 每句話沒有 pad 過的長度\n",
        "        \"prev_output_tokens\": prev_output_tokens, # 上面提到右 shift 一格後的目標序列\n",
        "    },\n",
        "    \"target\": target, # 目標序列\n",
        "}\n",
        "```"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "FzHJ6yIBsqk3"
      },
      "source": [
        "# 定義模型架構\n",
        "* 我們一樣繼承 fairseq 的 encoder, decoder 和 model, 這樣測試階段才能直接用他寫好的 beam search 函式"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "i4MzO2_csqk3"
      },
      "source": [
        "from fairseq.models import (\n",
        "    FairseqEncoder, \n",
        "    FairseqIncrementalDecoder,\n",
        "    FairseqEncoderDecoderModel\n",
        ")"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "FZb9bQbDsqk4"
      },
      "source": [
        "## Encoder 編碼器"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "6OuABfCXsqk4"
      },
      "source": [
        "- seq2seq 模型的編碼器為 RNN 或 Transformer Encoder，以下說明以 RNN 為例，Transformer 略有不同。對於每個輸入，Encoder 會輸出一個向量和一個隱藏狀態(hidden state)，並將隱藏狀態用於下一個輸入。換句話說，Encoder 會逐步讀取輸入序列，並在每個 timestep 輸出單個向量，以及在最後 timestep 輸出最終隱藏狀態(content vector)\n",
        "- 參數:\n",
        "  - *args*\n",
        "      - encoder_embed_dim 是 embedding 的維度，主要將 one-hot vector 的單詞向量壓縮到指定的維度，主要是為了降維和濃縮資訊的功用\n",
        "      - encoder_ffn_embed_dim 是 RNN 輸出和隱藏狀態的維度(hidden dimension)\n",
        "      - encoder_layers 是 RNN 要疊多少層\n",
        "      - dropout 是決定有多少的機率會將某個節點變為 0，主要是為了防止 overfitting ，一般來說是在訓練時使用，測試時則不使用\n",
        "  - *dictionary*: fairseq 幫我們做好的 dictionary. 在此用來得到 padding index，好用來得到 encoder padding mask. \n",
        "  - *embed_tokens*: 事先做好的詞嵌入 (nn.Embedding)\n",
        "\n",
        "- 輸入: \n",
        "    - *src_tokens*: 英文的整數序列 e.g. 1, 28, 29, 205, 2 \n",
        "- 輸出: \n",
        "    - *outputs*: 最上層 RNN 每個 timestep 的輸出，後續可以用 Attention 再進行處理\n",
        "    - *final_hiddens*: 每層最終 timestep 的隱藏狀態，將傳遞到 Decoder 進行解碼\n",
        "    - *encoder_padding_mask*: 告訴我們哪些是位置的資訊不重要。\n"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "aN5JjOzYsqk5"
      },
      "source": [
        "class RNNEncoder(FairseqEncoder):\n",
        "    def __init__(self, args, dictionary, embed_tokens):\n",
        "        super().__init__(dictionary)\n",
        "        self.embed_tokens = embed_tokens\n",
        "        \n",
        "        self.embed_dim = args.encoder_embed_dim\n",
        "        self.hidden_dim = args.encoder_ffn_embed_dim\n",
        "        self.num_layers = args.encoder_layers\n",
        "        \n",
        "        self.dropout_in_module = nn.Dropout(args.dropout)\n",
        "        self.rnn = nn.GRU(\n",
        "            self.embed_dim, \n",
        "            self.hidden_dim, \n",
        "            self.num_layers, \n",
        "            dropout=args.dropout, \n",
        "            batch_first=False, \n",
        "            bidirectional=True\n",
        "        )\n",
        "        self.dropout_out_module = nn.Dropout(args.dropout)\n",
        "        \n",
        "        self.padding_idx = dictionary.pad()\n",
        "        \n",
        "    def combine_bidir(self, outs, bsz: int):\n",
        "        out = outs.view(self.num_layers, 2, bsz, -1).transpose(1, 2).contiguous()\n",
        "        return out.view(self.num_layers, bsz, -1)\n",
        "\n",
        "    def forward(self, src_tokens, **unused):\n",
        "        bsz, seqlen = src_tokens.size()\n",
        "        \n",
        "        # get embeddings\n",
        "        x = self.embed_tokens(src_tokens)\n",
        "        x = self.dropout_in_module(x)\n",
        "\n",
        "        # B x T x C -> T x B x C\n",
        "        x = x.transpose(0, 1)\n",
        "        \n",
        "        # 過雙向RNN\n",
        "        h0 = x.new_zeros(2 * self.num_layers, bsz, self.hidden_dim)\n",
        "        x, final_hiddens = self.rnn(x, h0)\n",
        "        outputs = self.dropout_out_module(x)\n",
        "        # outputs = [sequence len, batch size, hid dim * directions] 是最上層RNN的輸出\n",
        "        # hidden =  [num_layers * directions, batch size  , hid dim]\n",
        "        \n",
        "        # 因為 Encoder 是雙向的RNN，所以需要將同一層兩個方向的 hidden state 接在一起\n",
        "        final_hiddens = self.combine_bidir(final_hiddens, bsz)\n",
        "        # hidden =  [num_layers x batch x num_directions*hidden]\n",
        "        \n",
        "        encoder_padding_mask = src_tokens.eq(self.padding_idx).t()\n",
        "        return tuple(\n",
        "            (\n",
        "                outputs,  # seq_len x batch x hidden\n",
        "                final_hiddens,  # num_layers x batch x num_directions*hidden\n",
        "                encoder_padding_mask,  # seq_len x batch\n",
        "            )\n",
        "        )\n",
        "    \n",
        "    def reorder_encoder_out(self, encoder_out, new_order):\n",
        "        # 這個beam search時會用到，意義並不是很重要\n",
        "        return tuple(\n",
        "            (\n",
        "                encoder_out[0].index_select(1, new_order),\n",
        "                encoder_out[1].index_select(1, new_order),\n",
        "                encoder_out[2].index_select(1, new_order),\n",
        "            )\n",
        "        )"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "8XVwBreMsqk5"
      },
      "source": [
        "## Attention"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "uVS1ovJmsqk5"
      },
      "source": [
        "- 當輸入過長，或是單獨靠 “content vector” 無法取得整個輸入的意思時，用 Attention Mechanism 來提供 Decoder 更多的資訊\n",
        "- 根據現在 **Decoder embeddings** ，去計算在 **Encoder outputs** 中，那些與其有較高的關係，根據關係的數值來把 Encoder outputs 平均起來作為 **Decoder** RNN 的輸入 \n",
        "- 常見 Attention 的實作是用 Neural Network / Dot Product 來算 **query** (decoder embeddings) 和 **key** (Encoder outputs) 之間的關係，再對所有算出來的數值做 **softmax** 得到分佈，最後根據這個分佈對 **values** (Encoder outputs) 做 **weight sum**\n",
        "\n",
        "- 參數:\n",
        "  - *input_embed_dim*: key 的維度，應是 decoder 要做 attend 時的向量的維度\n",
        "  - *source_embed_dim*: query 的維度，應是要被 attend 的向量(encoder outputs)的維度\n",
        "  - *output_embed_dim*: value 的維度，應是做完 attention 後，下一層預期的向量維度\n",
        "\n",
        "- 輸入: \n",
        "    - *inputs*: 就是 key，要 attend 別人的向量\n",
        "    - *encoder_outputs*: 是 query/value，被 attend 的向量\n",
        "    - *encoder_padding_mask*: 告訴我們哪些是位置的資訊不重要。\n",
        "- 輸出: \n",
        "    - *output*: 做完 attention 後的 context vector\n",
        "    - *attention score*: attention 的分布\n"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "R0_f9UjHsqk6"
      },
      "source": [
        "class AttentionLayer(nn.Module):\n",
        "    def __init__(self, input_embed_dim, source_embed_dim, output_embed_dim, bias=False):\n",
        "        super().__init__()\n",
        "\n",
        "        self.input_proj = nn.Linear(input_embed_dim, source_embed_dim, bias=bias)\n",
        "        self.output_proj = nn.Linear(\n",
        "            input_embed_dim + source_embed_dim, output_embed_dim, bias=bias\n",
        "        )\n",
        "\n",
        "    def forward(self, inputs, encoder_outputs, encoder_padding_mask):\n",
        "        # inputs: T, B, dim\n",
        "        # encoder_outputs: S x B x dim\n",
        "        # padding mask:  S x B\n",
        "        \n",
        "        # convert all to batch first\n",
        "        inputs = inputs.transpose(1,0) # B, T, dim\n",
        "        encoder_outputs = encoder_outputs.transpose(1,0) # B, S, dim\n",
        "        encoder_padding_mask = encoder_padding_mask.transpose(1,0) # B, S\n",
        "        \n",
        "        # 投影到encoder_outputs的維度\n",
        "        x = self.input_proj(inputs)\n",
        "\n",
        "        # 計算attention\n",
        "        # (B, T, dim) x (B, dim, S) = (B, T, S)\n",
        "        attn_scores = torch.bmm(x, encoder_outputs.transpose(1,2))\n",
        "\n",
        "        # 擋住padding位置的attention\n",
        "        if encoder_padding_mask is not None:\n",
        "            # 利用broadcast  B, S -> (B, 1, S)\n",
        "            encoder_padding_mask = encoder_padding_mask.unsqueeze(1)\n",
        "            attn_scores = (\n",
        "                attn_scores.float()\n",
        "                .masked_fill_(encoder_padding_mask, float(\"-inf\"))\n",
        "                .type_as(attn_scores)\n",
        "            )  # FP16 support: cast to float and back\n",
        "\n",
        "        # 在source對應維度softmax\n",
        "        attn_scores = F.softmax(attn_scores, dim=-1)\n",
        "\n",
        "        # 形狀 (B, T, S) x (B, S, dim) = (B, T, dim) 加權平均\n",
        "        x = torch.bmm(attn_scores, encoder_outputs)\n",
        "\n",
        "        # (B, T, dim)\n",
        "        x = torch.cat((x, inputs), dim=-1)\n",
        "        x = torch.tanh(self.output_proj(x)) # concat + linear + tanh\n",
        "        \n",
        "        # 回復形狀 (B, T, dim) -> (T, B, dim)\n",
        "        return x.transpose(1,0), attn_scores"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "PGCeUiNzsqk6"
      },
      "source": [
        "## Decoder 解碼器"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "HCUdUREHsqk6"
      },
      "source": [
        "* 解碼器的 hidden states 會用編碼器最終隱藏狀態來初始化(content vector)\n",
        "* 解碼器同時也根據目前 timestep 的輸入(也就是前幾個 timestep 的 output)，改變 hidden states，並輸出結果 \n",
        "* 如果加入 attention 可以使表現更好\n",
        "* 我們把 seq2seq 步驟寫在解碼器裡，好讓等等 Seq2Seq 這個型別可以通用 RNN 和 Transformer，而不用再改寫\n",
        "- 參數:\n",
        "  - *args*\n",
        "      - decoder_embed_dim 是解碼器 embedding 的維度，類同 encoder_embed_dim，\n",
        "      - decoder_ffn_embed_dim 是解碼器 RNN 的隱藏維度，類同 encoder_ffn_embed_dim\n",
        "      - decoder_layers 解碼器 RNN 的層數\n",
        "      - share_decoder_input_output_embed 通常 decoder 最後輸出的投影矩陣會和輸入 embedding 共用參數\n",
        "  - *dictionary*: fairseq 幫我們做好的 dictionary.\n",
        "  - *embed_tokens*: 事先做好的詞嵌入(nn.Embedding)\n",
        "- 輸入: \n",
        "    - *prev_output_tokens*: 英文的整數序列 e.g. 1, 28, 29, 205, 2 已經 shift 一格的 target\n",
        "    - *encoder_out*: 編碼器的輸出\n",
        "    - *incremental_state*: 這是測試階段為了加速，所以會記錄每個 timestep 的 hidden state 詳見 forward\n",
        "- 輸出: \n",
        "    - *outputs*: decoder 每個 timestep 的 logits，還沒經過 softmax 的分布\n",
        "    - *extra*: 沒用到"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "vIG203d9sqk7"
      },
      "source": [
        "class RNNDecoder(FairseqIncrementalDecoder):\n",
        "    def __init__(self, args, dictionary, embed_tokens):\n",
        "        super().__init__(dictionary)\n",
        "        self.embed_tokens = embed_tokens\n",
        "        \n",
        "        assert args.decoder_layers == args.encoder_layers, f\"\"\"seq2seq rnn requires that encoder \n",
        "        and decoder have same layers of rnn. got: {args.encoder_layers, args.decoder_layers}\"\"\"\n",
        "        assert args.decoder_ffn_embed_dim == args.encoder_ffn_embed_dim*2, f\"\"\"seq2seq-rnn requires \n",
        "        that decoder hidden to be 2*encoder hidden dim. got: {args.decoder_ffn_embed_dim, args.encoder_ffn_embed_dim*2}\"\"\"\n",
        "        \n",
        "        self.embed_dim = args.decoder_embed_dim\n",
        "        self.hidden_dim = args.decoder_ffn_embed_dim\n",
        "        self.num_layers = args.decoder_layers\n",
        "        \n",
        "        \n",
        "        self.dropout_in_module = nn.Dropout(args.dropout)\n",
        "        self.rnn = nn.GRU(\n",
        "            self.embed_dim, \n",
        "            self.hidden_dim, \n",
        "            self.num_layers, \n",
        "            dropout=args.dropout, \n",
        "            batch_first=False, \n",
        "            bidirectional=False\n",
        "        )\n",
        "        self.attention = AttentionLayer(\n",
        "            self.embed_dim, self.hidden_dim, self.embed_dim, bias=False\n",
        "        ) \n",
        "        # self.attention = None\n",
        "        self.dropout_out_module = nn.Dropout(args.dropout)\n",
        "        \n",
        "        if self.hidden_dim != self.embed_dim:\n",
        "            self.project_out_dim = nn.Linear(self.hidden_dim, self.embed_dim)\n",
        "        else:\n",
        "            self.project_out_dim = None\n",
        "        \n",
        "        if args.share_decoder_input_output_embed:\n",
        "            self.output_projection = nn.Linear(\n",
        "                self.embed_tokens.weight.shape[1],\n",
        "                self.embed_tokens.weight.shape[0],\n",
        "                bias=False,\n",
        "            )\n",
        "            self.output_projection.weight = self.embed_tokens.weight\n",
        "        else:\n",
        "            self.output_projection = nn.Linear(\n",
        "                self.output_embed_dim, len(dictionary), bias=False\n",
        "            )\n",
        "            nn.init.normal_(\n",
        "                self.output_projection.weight, mean=0, std=self.output_embed_dim ** -0.5\n",
        "            )\n",
        "        \n",
        "    def forward(self, prev_output_tokens, encoder_out, incremental_state=None, **unused):\n",
        "        # 取出encoder的輸出\n",
        "        encoder_outputs, encoder_hiddens, encoder_padding_mask = encoder_out\n",
        "        # outputs:          seq_len x batch x num_directions*hidden\n",
        "        # encoder_hiddens:  num_layers x batch x num_directions*encoder_hidden\n",
        "        # padding_mask:     seq_len x batch\n",
        "        \n",
        "        if incremental_state is not None and len(incremental_state) > 0:\n",
        "            # 有上個timestep留下的資訊，讀進來就可以繼續decode，不用從bos重來\n",
        "            prev_output_tokens = prev_output_tokens[:, -1:]\n",
        "            cache_state = self.get_incremental_state(incremental_state, \"cached_state\")\n",
        "            prev_hiddens = cache_state[\"prev_hiddens\"]\n",
        "        else:\n",
        "            # 沒有incremental state代表這是training或者是test time時的第一步\n",
        "            # 準備seq2seq: 把encoder_hiddens pass進去decoder的hidden states\n",
        "            prev_hiddens = encoder_hiddens\n",
        "        \n",
        "        bsz, seqlen = prev_output_tokens.size()\n",
        "        \n",
        "        # embed tokens\n",
        "        x = self.embed_tokens(prev_output_tokens)\n",
        "        x = self.dropout_in_module(x)\n",
        "\n",
        "        # B x T x C -> T x B x C\n",
        "        x = x.transpose(0, 1)\n",
        "                \n",
        "        # 做decoder-to-encoder attention\n",
        "        if self.attention is not None:\n",
        "            x, attn = self.attention(x, encoder_outputs, encoder_padding_mask)\n",
        "                        \n",
        "        # 過單向RNN\n",
        "        x, final_hiddens = self.rnn(x, prev_hiddens)\n",
        "        # outputs = [sequence len, batch size, hid dim]\n",
        "        # hidden =  [num_layers * directions, batch size  , hid dim]\n",
        "        x = self.dropout_out_module(x)\n",
        "                \n",
        "        # 投影到embedding size (如果hidden 和embed size不一樣，然後share_embedding又設成True,需要額外project一次)\n",
        "        if self.project_out_dim != None:\n",
        "            x = self.project_out_dim(x)\n",
        "        \n",
        "        # 投影到vocab size 的分佈\n",
        "        x = self.output_projection(x)\n",
        "        \n",
        "        # T x B x C -> B x T x C\n",
        "        x = x.transpose(1, 0)\n",
        "        \n",
        "        # 如果是Incremental, 記錄這個timestep的hidden states, 下個timestep讀回來\n",
        "        cache_state = {\n",
        "            \"prev_hiddens\": final_hiddens,\n",
        "        }\n",
        "        self.set_incremental_state(incremental_state, \"cached_state\", cache_state)\n",
        "        \n",
        "        return x, None\n",
        "    \n",
        "    def reorder_incremental_state(\n",
        "        self,\n",
        "        incremental_state,\n",
        "        new_order,\n",
        "    ):\n",
        "        # 這個beam search時會用到，意義並不是很重要\n",
        "        cache_state = self.get_incremental_state(incremental_state, \"cached_state\")\n",
        "        prev_hiddens = cache_state[\"prev_hiddens\"]\n",
        "        prev_hiddens = [p.index_select(0, new_order) for p in prev_hiddens]\n",
        "        cache_state = {\n",
        "            \"prev_hiddens\": torch.stack(prev_hiddens),\n",
        "        }\n",
        "        self.set_incremental_state(incremental_state, \"cached_state\", cache_state)\n",
        "        return"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "y55fockAsqk7"
      },
      "source": [
        "## Seq2Seq\n",
        "- 由 **Encoder** 和 **Decoder** 組成\n",
        "- 接收輸入並傳給 **Encoder** \n",
        "- 將 **Encoder** 的輸出傳給 **Decoder**\n",
        "- **Decoder** 根據前幾個 timestep 的輸出和 **Encoder** 輸出進行解碼  \n",
        "- 當解碼完成後，將 **Decoder** 的輸出傳回 "
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "P07--YsJsqk7"
      },
      "source": [
        "class Seq2Seq(FairseqEncoderDecoderModel):\n",
        "    def __init__(self, args, encoder, decoder):\n",
        "        super().__init__(encoder, decoder)\n",
        "        self.args = args\n",
        "    \n",
        "    def forward(\n",
        "        self,\n",
        "        src_tokens,\n",
        "        src_lengths,\n",
        "        prev_output_tokens,\n",
        "        return_all_hiddens: bool = True,\n",
        "    ):\n",
        "        \"\"\"\n",
        "        Run the forward pass for an encoder-decoder model.\n",
        "        \"\"\"\n",
        "        encoder_out = self.encoder(\n",
        "            src_tokens, src_lengths=src_lengths, return_all_hiddens=return_all_hiddens\n",
        "        )\n",
        "        logits, extra = self.decoder(\n",
        "            prev_output_tokens,\n",
        "            encoder_out=encoder_out,\n",
        "            src_lengths=src_lengths,\n",
        "            return_all_hiddens=return_all_hiddens,\n",
        "        )\n",
        "        return logits, extra"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "nSggF_v8sqk8"
      },
      "source": [
        "# 模型初始化"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "kxI2oy9Fsqk8"
      },
      "source": [
        "# # HINT: transformer 架構\n",
        "# from fairseq.models.transformer import (\n",
        "#     TransformerEncoder, \n",
        "#     TransformerDecoder,\n",
        "# )\n",
        "\n",
        "def build_model(args, task):\n",
        "    \"\"\" 按照參數設定建置模型 \"\"\"\n",
        "    src_dict, tgt_dict = task.source_dictionary, task.target_dictionary\n",
        "\n",
        "    # 詞嵌入\n",
        "    encoder_embed_tokens = nn.Embedding(len(src_dict), args.encoder_embed_dim, src_dict.pad())\n",
        "    decoder_embed_tokens = nn.Embedding(len(tgt_dict), args.decoder_embed_dim, tgt_dict.pad())\n",
        "    \n",
        "    # 編碼器與解碼器\n",
        "    # TODO: 替換成 TransformerEncoder 和 TransformerDecoder\n",
        "    encoder = RNNEncoder(args, src_dict, encoder_embed_tokens)\n",
        "    decoder = RNNDecoder(args, tgt_dict, decoder_embed_tokens)\n",
        "    \n",
        "    # 序列到序列模型\n",
        "    model = Seq2Seq(args, encoder, decoder)\n",
        "    \n",
        "    # 序列到序列模型的初始化很重要 需要特別處理\n",
        "    def init_params(module):\n",
        "        from fairseq.modules import MultiheadAttention\n",
        "        if isinstance(module, nn.Linear):\n",
        "            module.weight.data.normal_(mean=0.0, std=0.02)\n",
        "            if module.bias is not None:\n",
        "                module.bias.data.zero_()\n",
        "        if isinstance(module, nn.Embedding):\n",
        "            module.weight.data.normal_(mean=0.0, std=0.02)\n",
        "            if module.padding_idx is not None:\n",
        "                module.weight.data[module.padding_idx].zero_()\n",
        "        if isinstance(module, MultiheadAttention):\n",
        "            module.q_proj.weight.data.normal_(mean=0.0, std=0.02)\n",
        "            module.k_proj.weight.data.normal_(mean=0.0, std=0.02)\n",
        "            module.v_proj.weight.data.normal_(mean=0.0, std=0.02)\n",
        "        if isinstance(module, nn.RNNBase):\n",
        "            for name, param in module.named_parameters():\n",
        "                if \"weight\" in name or \"bias\" in name:\n",
        "                    param.data.uniform_(-0.1, 0.1)\n",
        "            \n",
        "    # 初始化模型\n",
        "    model.apply(init_params)\n",
        "    return model"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "6HZDLRf2sqk_"
      },
      "source": [
        "## 設定模型相關參數\n",
        "參考參數\n",
        "\n",
        "|model|embedding dim|encoder ffn|encoder layers|decoder ffn|decoder layers|\n",
        "|-|-|-|-|-|-|\n",
        "|RNN|256|512|1|1024|1|\n",
        "|Transformer|256|1024|4|1024|4|\n",
        "\n",
        "Strong baseline 用的參數可以參考 [Attention is all you need](#vaswani2017) 的 Table 3 的 transformer-base"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "Nrmks6QcsqlA"
      },
      "source": [
        "arch_args = Namespace(\n",
        "    encoder_embed_dim=256,\n",
        "    encoder_ffn_embed_dim=512,\n",
        "    encoder_layers=1,\n",
        "    decoder_embed_dim=256,\n",
        "    decoder_ffn_embed_dim=1024,\n",
        "    decoder_layers=1,\n",
        "    share_decoder_input_output_embed=True,\n",
        "    dropout=0.3,\n",
        ")\n",
        "\n",
        "# # HINT: 補上Transformer用的參數\n",
        "# def add_transformer_args(args):\n",
        "#     args.encoder_attention_heads=4\n",
        "#     args.encoder_normalize_before=True\n",
        "    \n",
        "#     args.decoder_attention_heads=4\n",
        "#     args.decoder_normalize_before=True\n",
        "    \n",
        "#     args.activation_fn=\"relu\"\n",
        "#     args.max_source_positions=1024\n",
        "#     args.max_target_positions=1024\n",
        "    \n",
        "#     # 補上我們沒有設定的Transformer預設參數\n",
        "#     from fairseq.models.transformer import base_architecture \n",
        "#     base_architecture(arch_args)\n",
        "\n",
        "# add_transformer_args(arch_args)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "nPtuo_h5sqlA"
      },
      "source": [
        "if config.use_wandb:\n",
        "    wandb.config.update(vars(arch_args))"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "KshsZuGisqlB"
      },
      "source": [
        "model = build_model(arch_args, task)\n",
        "logger.info(model)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "zUmk-a8tsqlB"
      },
      "source": [
        "# Optimization 最佳化"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "yvqT93nNsqlB"
      },
      "source": [
        "## Loss: Label Smoothing Regularization\n",
        "* 讓模型學習輸出較不集中的分佈，防止模型過度自信\n",
        "* 有時候Ground Truth並非唯一答案，所以在算loss時，我們會保留一部份機率給正確答案以外的label\n",
        "* 可以有效防止過度擬合\n",
        "\n",
        "code [source](https://fairseq.readthedocs.io/en/latest/_modules/fairseq/criterions/label_smoothed_cross_entropy.html)"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "TKjsrlJCsqlC"
      },
      "source": [
        "class LabelSmoothedCrossEntropyCriterion(nn.Module):\n",
        "    def __init__(self, smoothing, ignore_index=None, reduce=True):\n",
        "        super().__init__()\n",
        "        self.smoothing = smoothing\n",
        "        self.ignore_index = ignore_index\n",
        "        self.reduce = reduce\n",
        "    \n",
        "    def forward(self, lprobs, target):\n",
        "        if target.dim() == lprobs.dim() - 1:\n",
        "            target = target.unsqueeze(-1)\n",
        "        # nll: Negative log likelihood，當目標是one-hot時的cross-entropy loss. 以下同 F.nll_loss\n",
        "        nll_loss = -lprobs.gather(dim=-1, index=target)\n",
        "        # 將一部分正確答案的機率分配給其他label 所以當計算cross-entropy時等於把所有label的log prob加起來\n",
        "        smooth_loss = -lprobs.sum(dim=-1, keepdim=True)\n",
        "        if self.ignore_index is not None:\n",
        "            pad_mask = target.eq(self.ignore_index)\n",
        "            nll_loss.masked_fill_(pad_mask, 0.0)\n",
        "            smooth_loss.masked_fill_(pad_mask, 0.0)\n",
        "        else:\n",
        "            nll_loss = nll_loss.squeeze(-1)\n",
        "            smooth_loss = smooth_loss.squeeze(-1)\n",
        "        if self.reduce:\n",
        "            nll_loss = nll_loss.sum()\n",
        "            smooth_loss = smooth_loss.sum()\n",
        "        # 計算cross-entropy時 加入分配給其他label的loss\n",
        "        eps_i = self.smoothing / lprobs.size(-1)\n",
        "        loss = (1.0 - self.smoothing) * nll_loss + eps_i * smooth_loss\n",
        "        return loss\n",
        "\n",
        "# 一般都用0.1效果就很好了\n",
        "criterion = LabelSmoothedCrossEntropyCriterion(\n",
        "    smoothing=0.1,\n",
        "    ignore_index=task.target_dictionary.pad(),\n",
        ")"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "5BsgL03ssqlC"
      },
      "source": [
        "## Optimizer: Adam + lr scheduling\n",
        "Inverse square root 排程對於訓練 Transformer 時的穩定性很重要，後來也用在 RNN 上。\n",
        "根據底下公式來更新 learning rate，前期線性增長，後期根據更新步數方根的倒數來遞減。\n",
        "$$lrate = d_{\\text{model}}^{-0.5}\\cdot\\min({step\\_num}^{-0.5},{step\\_num}\\cdot{warmup\\_steps}^{-1.5})$$\n",
        "code [source](https://nlp.seas.harvard.edu/2018/04/03/attention.html)"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "n6itJcIDsqlC"
      },
      "source": [
        "class NoamOpt:\n",
        "    \"Optim wrapper that implements rate.\"\n",
        "    def __init__(self, model_size, factor, warmup, optimizer):\n",
        "        self.optimizer = optimizer\n",
        "        self._step = 0\n",
        "        self.warmup = warmup\n",
        "        self.factor = factor\n",
        "        self.model_size = model_size\n",
        "        self._rate = 0\n",
        "    \n",
        "    @property\n",
        "    def param_groups(self):\n",
        "        return self.optimizer.param_groups\n",
        "        \n",
        "    def multiply_grads(self, c):\n",
        "        \"\"\"Multiplies grads by a constant *c*.\"\"\"                \n",
        "        for group in self.param_groups:\n",
        "            for p in group['params']:\n",
        "                if p.grad is not None:\n",
        "                    p.grad.data.mul_(c)\n",
        "        \n",
        "    def step(self):\n",
        "        \"Update parameters and rate\"\n",
        "        self._step += 1\n",
        "        rate = self.rate()\n",
        "        for p in self.param_groups:\n",
        "            p['lr'] = rate\n",
        "        self._rate = rate\n",
        "        self.optimizer.step()\n",
        "        \n",
        "    def rate(self, step = None):\n",
        "        \"Implement `lrate` above\"\n",
        "        if step is None:\n",
        "            step = self._step\n",
        "        return 0 if not step else self.factor * \\\n",
        "            (self.model_size ** (-0.5) *\n",
        "            min(step ** (-0.5), step * self.warmup ** (-1.5)))"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "vbR_x6mesqlC"
      },
      "source": [
        "## 排程視覺化"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "4pH_JrP9sqlD"
      },
      "source": [
        "optimizer = NoamOpt(\n",
        "    model_size=arch_args.encoder_embed_dim, \n",
        "    factor=config.lr_factor, \n",
        "    warmup=config.lr_warmup, \n",
        "    optimizer=torch.optim.AdamW(model.parameters(), lr=0, betas=(0.9, 0.98), eps=1e-9, weight_decay=0.0001))\n",
        "plt.plot(np.arange(1, 100000), [optimizer.rate(i) for i in range(1, 100000)])\n",
        "plt.legend([f\"{optimizer.model_size}:{optimizer.warmup}\"])\n",
        "None"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "BVLxbyMjsqlD"
      },
      "source": [
        "# 訓練步驟"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "u6-c3sxPsqlD"
      },
      "source": [
        "## Training 訓練"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "1ELSNqtdsqlD"
      },
      "source": [
        "from fairseq.data import iterators\n",
        "from torch.cuda.amp import GradScaler, autocast\n",
        "\n",
        "def train_one_epoch(epoch_itr, model, task, criterion, optimizer, accum_steps=1):\n",
        "    itr = epoch_itr.next_epoch_itr(shuffle=True)\n",
        "    itr = iterators.GroupedIterator(itr, accum_steps) # 梯度累積: 每 accum_steps 個 sample 更新一次\n",
        "    \n",
        "    stats = {\"loss\": []}\n",
        "    scaler = GradScaler() # 混和精度訓練 automatic mixed precision (amp) \n",
        "    \n",
        "    model.train()\n",
        "    progress = tqdm.tqdm(itr, desc=f\"train epoch {epoch_itr.epoch}\", leave=False)\n",
        "    for samples in progress:\n",
        "        model.zero_grad()\n",
        "        accum_loss = 0\n",
        "        sample_size = 0\n",
        "        # 梯度累積: 每 accum_steps 個 sample 更新一次\n",
        "        for i, sample in enumerate(samples):\n",
        "            if i == 1:\n",
        "                # emptying the CUDA cache after the first step can reduce the chance of OOM\n",
        "                torch.cuda.empty_cache()\n",
        "\n",
        "            sample = utils.move_to_cuda(sample, device=device)\n",
        "            target = sample[\"target\"]\n",
        "            sample_size_i = sample[\"ntokens\"]\n",
        "            sample_size += sample_size_i\n",
        "            \n",
        "            # 混和精度訓練 \n",
        "            with autocast():\n",
        "                net_output = model.forward(**sample[\"net_input\"])\n",
        "                lprobs = F.log_softmax(net_output[0], -1)            \n",
        "                loss = criterion(lprobs.view(-1, lprobs.size(-1)), target.view(-1))\n",
        "                \n",
        "                # logging\n",
        "                accum_loss += loss.item()\n",
        "                # back-prop\n",
        "                scaler.scale(loss).backward()                \n",
        "        \n",
        "        scaler.unscale_(optimizer)\n",
        "        optimizer.multiply_grads(1 / (sample_size or 1.0)) # (sample_size or 1.0) handles the case of a zero gradient\n",
        "        gnorm = nn.utils.clip_grad_norm_(model.parameters(), config.clip_norm) # 梯度裁剪 防止梯度爆炸\n",
        "        \n",
        "        scaler.step(optimizer)\n",
        "        scaler.update()\n",
        "        \n",
        "        # logging\n",
        "        loss_print = accum_loss/sample_size\n",
        "        stats[\"loss\"].append(loss_print)\n",
        "        progress.set_postfix(loss=loss_print)\n",
        "        if config.use_wandb:\n",
        "            wandb.log({\n",
        "                \"train/loss\": loss_print,\n",
        "                \"train/grad_norm\": gnorm.item(),\n",
        "                \"train/lr\": optimizer.rate(),\n",
        "                \"train/sample_size\": sample_size,\n",
        "            })\n",
        "        \n",
        "    loss_print = np.mean(stats[\"loss\"])\n",
        "    logger.info(f\"training loss: {loss_print:.4f}\")\n",
        "    return stats"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "w0Tc_VhxsqlE"
      },
      "source": [
        "## Validation & Inference 檢驗和推論\n",
        "為防止訓練發生過度擬合，每過一段時間要做一次檢測，計算模型在未看過的資料上的表現。\n",
        "- 過程基本上和training一樣，另外加上 inference\n",
        "- 檢驗完畢可順便儲存模型參數\n",
        "\n",
        "單看 validation loss，我們很難知道模型真實的效能\n",
        "- 直接用當前模型去生成翻譯結果 (hypothesis)，再和正確答案 (reference) 計算 BLEU score\n",
        "- 也可用肉眼看翻譯結果的好壞\n",
        "- 我們用 fairseq 寫好的 sequence generator 來進行 beam search 生成翻譯結果"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "bvysFsA7sqlE"
      },
      "source": [
        "# fairseq 的 beam search generator\n",
        "# 給定模型和輸入序列，用 beam search 生成翻譯結果\n",
        "sequence_generator = task.build_generator([model], config)\n",
        "\n",
        "def decode(toks, dictionary):\n",
        "    # 從 Tensor 轉成人看得懂的句子\n",
        "    s = dictionary.string(\n",
        "        toks.int().cpu(),\n",
        "        config.post_process,\n",
        "    )\n",
        "    return s if s else \"<unk>\"\n",
        "\n",
        "def inference_step(sample, model):\n",
        "    gen_out = sequence_generator.generate([model], sample)\n",
        "    srcs = []\n",
        "    hyps = []\n",
        "    refs = []\n",
        "    for i in range(len(gen_out)):\n",
        "        # 對於每個 sample, 收集輸入，輸出和參考答案，稍後計算 BLEU\n",
        "        srcs.append(decode(\n",
        "            utils.strip_pad(sample[\"net_input\"][\"src_tokens\"][i], task.source_dictionary.pad()), \n",
        "            task.source_dictionary,\n",
        "        ))\n",
        "        hyps.append(decode(\n",
        "            gen_out[i][0][\"tokens\"], # 0 代表取出 beam 內分數第一的輸出結果\n",
        "            task.target_dictionary,\n",
        "        ))\n",
        "        refs.append(decode(\n",
        "            utils.strip_pad(sample[\"target\"][i], task.target_dictionary.pad()), \n",
        "            task.target_dictionary,\n",
        "        ))\n",
        "    return srcs, hyps, refs"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "xesFn8vWsqlE"
      },
      "source": [
        "import shutil\n",
        "import sacrebleu\n",
        "\n",
        "def validate(model, task, criterion, log_to_wandb=True):\n",
        "    logger.info('begin validation')\n",
        "    itr = load_data_iterator(task, \"valid\", 1, config.max_tokens, config.num_workers).next_epoch_itr(shuffle=False)\n",
        "    \n",
        "    stats = {\"loss\":[], \"bleu\": 0, \"srcs\":[], \"hyps\":[], \"refs\":[]}\n",
        "    srcs = []\n",
        "    hyps = []\n",
        "    refs = []\n",
        "    \n",
        "    model.eval()\n",
        "    progress = tqdm.tqdm(itr, desc=f\"validation\", leave=False)\n",
        "    with torch.no_grad():\n",
        "        for i, sample in enumerate(progress):\n",
        "            # validation loss\n",
        "            sample = utils.move_to_cuda(sample, device=device)\n",
        "            net_output = model.forward(**sample[\"net_input\"])\n",
        "\n",
        "            lprobs = F.log_softmax(net_output[0], -1)\n",
        "            target = sample[\"target\"]\n",
        "            sample_size = sample[\"ntokens\"]\n",
        "            loss = criterion(lprobs.view(-1, lprobs.size(-1)), target.view(-1)) / sample_size\n",
        "            progress.set_postfix(valid_loss=loss.item())\n",
        "            stats[\"loss\"].append(loss)\n",
        "            \n",
        "            # 進行推論\n",
        "            s, h, r = inference_step(sample, model)\n",
        "            srcs.extend(s)\n",
        "            hyps.extend(h)\n",
        "            refs.extend(r)\n",
        "            \n",
        "    tok = 'zh' if task.cfg.target_lang == 'zh' else '13a'\n",
        "    stats[\"loss\"] = torch.stack(stats[\"loss\"]).mean().item()\n",
        "    stats[\"bleu\"] = sacrebleu.corpus_bleu(hyps, [refs], tokenize=tok) # 計算BLEU score\n",
        "    stats[\"srcs\"] = srcs\n",
        "    stats[\"hyps\"] = hyps\n",
        "    stats[\"refs\"] = refs\n",
        "    \n",
        "    if config.use_wandb and log_to_wandb:\n",
        "        wandb.log({\n",
        "            \"valid/loss\": stats[\"loss\"],\n",
        "            \"valid/bleu\": stats[\"bleu\"].score,\n",
        "        }, commit=False)\n",
        "    \n",
        "    showid = np.random.randint(len(hyps))\n",
        "    logger.info(\"example source: \" + srcs[showid])\n",
        "    logger.info(\"example hypothesis: \" + hyps[showid])\n",
        "    logger.info(\"example reference: \" + refs[showid])\n",
        "    \n",
        "    # show bleu results\n",
        "    logger.info(f\"validation loss:\\t{stats['loss']:.4f}\")\n",
        "    logger.info(stats[\"bleu\"].format())\n",
        "    return stats"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "5-Vql872sqlF"
      },
      "source": [
        "# 儲存及載入模型參數"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "hoi1rGy8sqlF"
      },
      "source": [
        "def validate_and_save(model, task, criterion, optimizer, epoch, save=True):   \n",
        "    stats = validate(model, task, criterion)\n",
        "    bleu = stats['bleu']\n",
        "    loss = stats['loss']\n",
        "    if save:\n",
        "        # save epoch checkpoints\n",
        "        savedir = Path(config.savedir).absolute()\n",
        "        savedir.mkdir(parents=True, exist_ok=True)\n",
        "        \n",
        "        check = {\n",
        "            \"model\": model.state_dict(),\n",
        "            \"stats\": {\"bleu\": bleu.score, \"loss\": loss},\n",
        "            \"optim\": {\"step\": optimizer._step}\n",
        "        }\n",
        "        torch.save(check, savedir/f\"checkpoint{epoch}.pt\")\n",
        "        shutil.copy(savedir/f\"checkpoint{epoch}.pt\", savedir/f\"checkpoint_last.pt\")\n",
        "        logger.info(f\"saved epoch checkpoint: {savedir}/checkpoint{epoch}.pt\")\n",
        "    \n",
        "        # save epoch samples\n",
        "        with open(savedir/f\"samples{epoch}.{config.source_lang}-{config.target_lang}.txt\", \"w\") as f:\n",
        "            for s, h in zip(stats[\"srcs\"], stats[\"hyps\"]):\n",
        "                f.write(f\"{s}\\t{h}\\n\")\n",
        "\n",
        "        # get best valid bleu    \n",
        "        if getattr(validate_and_save, \"best_bleu\", 0) < bleu.score:\n",
        "            validate_and_save.best_bleu = bleu.score\n",
        "            torch.save(check, savedir/f\"checkpoint_best.pt\")\n",
        "            \n",
        "        del_file = savedir / f\"checkpoint{epoch - config.keep_last_epochs}.pt\"\n",
        "        if del_file.exists():\n",
        "            del_file.unlink()\n",
        "    \n",
        "    return stats\n",
        "\n",
        "def try_load_checkpoint(model, optimizer=None, name=None):\n",
        "    name = name if name else \"checkpoint_last.pt\"\n",
        "    checkpath = Path(config.savedir)/name\n",
        "    if checkpath.exists():\n",
        "        check = torch.load(checkpath)\n",
        "        model.load_state_dict(check[\"model\"])\n",
        "        stats = check[\"stats\"]\n",
        "        step = \"unknown\"\n",
        "        if optimizer != None:\n",
        "            optimizer._step = step = check[\"optim\"][\"step\"]\n",
        "        logger.info(f\"loaded checkpoint {checkpath}: step={step} loss={stats['loss']} bleu={stats['bleu']}\")\n",
        "    else:\n",
        "        logger.info(f\"no checkpoints found at {checkpath}!\")"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "zuXYIch4sqlG"
      },
      "source": [
        "# 主程式\n",
        "## 訓練迴圈"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "8OAgZXIosqlG"
      },
      "source": [
        "model = model.to(device=device)\n",
        "criterion = criterion.to(device=device)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "a76jGLZdsqlG"
      },
      "source": [
        "!nvidia-smi"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "GlFkx_FbsqlG"
      },
      "source": [
        "logger.info(\"task: {}\".format(task.__class__.__name__))\n",
        "logger.info(\"encoder: {}\".format(model.encoder.__class__.__name__))\n",
        "logger.info(\"decoder: {}\".format(model.decoder.__class__.__name__))\n",
        "logger.info(\"criterion: {}\".format(criterion.__class__.__name__))\n",
        "logger.info(\"optimizer: {}\".format(optimizer.__class__.__name__))\n",
        "logger.info(\n",
        "    \"num. model params: {:,} (num. trained: {:,})\".format(\n",
        "        sum(p.numel() for p in model.parameters()),\n",
        "        sum(p.numel() for p in model.parameters() if p.requires_grad),\n",
        "    )\n",
        ")\n",
        "logger.info(f\"max tokens per batch = {config.max_tokens}, accumulate steps = {config.accum_steps}\")"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "scrolled": true,
        "id": "PyUw2n3csqlH"
      },
      "source": [
        "epoch_itr = load_data_iterator(task, \"train\", config.start_epoch, config.max_tokens, config.num_workers)\n",
        "try_load_checkpoint(model, optimizer, name=config.resume)\n",
        "while epoch_itr.next_epoch_idx <= config.max_epoch:\n",
        "    # train for one epoch\n",
        "    train_one_epoch(epoch_itr, model, task, criterion, optimizer, config.accum_steps)\n",
        "    stats = validate_and_save(model, task, criterion, optimizer, epoch=epoch_itr.epoch)\n",
        "    logger.info(\"end of epoch {}\".format(epoch_itr.epoch))    \n",
        "    epoch_itr = load_data_iterator(task, \"train\", epoch_itr.next_epoch_idx, config.max_tokens, config.num_workers)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "YdpMZTVFsqlH"
      },
      "source": [
        "# Submission 繳交檔案"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "4w89HjHksqlH"
      },
      "source": [
        "# 把幾個 checkpoint 平均起來可以達到 ensemble 的效果\n",
        "checkdir=config.savedir\n",
        "!python ./fairseq/scripts/average_checkpoints.py \\\n",
        "--inputs {checkdir} \\\n",
        "--num-epoch-checkpoints 5 \\\n",
        "--output {checkdir}/avg_last_5_checkpoint.pt"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "jXMFDPpZsqlH"
      },
      "source": [
        "## 確認生成繳交檔案的模型參數"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "TIPilIx5sqlI"
      },
      "source": [
        "# checkpoint_last.pt : 最後一次檢驗的檔案\n",
        "# checkpoint_best.pt : 檢驗 BLEU 最高的檔案\n",
        "# avg_last_5_checkpoint.pt:　最5後個檔案平均\n",
        "try_load_checkpoint(model, name=\"checkpoint_best.pt\")\n",
        "validate(model, task, criterion, log_to_wandb=False)\n",
        "None"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "fQe97WdQsqlI"
      },
      "source": [
        "## 進行預測"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "kgC-nfWVsqlI"
      },
      "source": [
        "def generate_prediction(model, task, split=\"test\", outfile=\"./prediction.txt\"):    \n",
        "    task.load_dataset(split=split, epoch=1)\n",
        "    itr = load_data_iterator(task, split, 1, config.max_tokens, config.num_workers).next_epoch_itr(shuffle=False)\n",
        "    \n",
        "    idxs = []\n",
        "    hyps = []\n",
        "\n",
        "    model.eval()\n",
        "    progress = tqdm.tqdm(itr, desc=f\"prediction\")\n",
        "    with torch.no_grad():\n",
        "        for i, sample in enumerate(progress):\n",
        "            # validation loss\n",
        "            sample = utils.move_to_cuda(sample, device=device)\n",
        "\n",
        "            # 進行推論\n",
        "            s, h, r = inference_step(sample, model)\n",
        "            \n",
        "            hyps.extend(h)\n",
        "            idxs.extend(list(sample['id']))\n",
        "            \n",
        "    # 根據 preprocess 時的順序排列\n",
        "    hyps = [x for _,x in sorted(zip(idxs,hyps))]\n",
        "    \n",
        "    with open(outfile, \"w\") as f:\n",
        "        for h in hyps:\n",
        "            f.write(h+\"\\n\")"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "8b1Y4DhMsqlI"
      },
      "source": [
        "generate_prediction(model, task)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "ZLQN9HsAsqlI"
      },
      "source": [
        "raise"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "FZcTQJXJsqlJ"
      },
      "source": [
        "# Back-translation"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "WWSkoh-RsqlJ"
      },
      "source": [
        "## 訓練一個反向的翻譯模型"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "E2Q4ceXAsqlJ"
      },
      "source": [
        "1. 將實驗的參數設定表中(config)的source_lang與target_lang互相交換\n",
        "2. 將實驗的參數設定表中(config)的savedir更改(ex. \"./checkpoints/rnn-back\")\n",
        "3. 訓練一個反向模型"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "xvuBOzKusqlJ"
      },
      "source": [
        "## 利用反向模型生成額外資料"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "JIn3EHgcsqlK"
      },
      "source": [
        "### 下載 monolingual data"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "jOkl5bahsqlK"
      },
      "source": [
        "mono_dataset_name = 'mono'"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "c8D7GbJ2sqlK"
      },
      "source": [
        "mono_prefix = Path(data_dir).absolute() / mono_dataset_name\n",
        "mono_prefix.mkdir(parents=True, exist_ok=True)\n",
        "\n",
        "urls = (\n",
        "    '\"https://onedrive.live.com/download?cid=3E549F3B24B238B4&resid=3E549F3B24B238B4%214986&authkey=AANUKbGfZx0kM80\"',\n",
        "# # If the above links die, use the following instead. \n",
        "#     \"https://www.csie.ntu.edu.tw/~r09922057/ML2021-hw5/ted_zh_corpus.deduped.gz\",\n",
        "# # If the above links die, use the following instead. \n",
        "#     \"https://mega.nz/#!vMNnDShR!4eHDxzlpzIpdpeQTD-htatU_C7QwcBTwGDaSeBqH534\",\n",
        ")\n",
        "file_names = (\n",
        "    'ted_zh_corpus.deduped.gz',\n",
        ")\n",
        "\n",
        "for u, f in zip(urls, file_names):\n",
        "    path = mono_prefix/f\n",
        "    if not path.exists():\n",
        "        if 'mega' in u:\n",
        "            !megadl {u} --path {path}\n",
        "        else:\n",
        "            !wget {u} -O {path}\n",
        "    else:\n",
        "        print(f'{f} is exist, skip downloading')\n",
        "    if path.suffix == \".tgz\":\n",
        "        !tar -xvf {path} -C {prefix}\n",
        "    elif path.suffix == \".zip\":\n",
        "        !unzip -o {path} -d {prefix}\n",
        "    elif path.suffix == \".gz\":\n",
        "        !gzip -fkd {path}"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "FHIFeljKsqlK"
      },
      "source": [
        "### TODO: 清理資料集\n",
        "\n",
        "1. 將太長、太短的句子移除\n",
        "2. 統一標點符號\n",
        "\n",
        "hint: 可以使用clean_s()來協助"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "7HYkgXXqsqlL"
      },
      "source": [
        ""
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "WKHaAcCcsqlL"
      },
      "source": [
        "### TODO: Subword Units\n",
        "\n",
        "用反向模型的 spm model 將資料切成 subword units\n",
        "\n",
        "hint: spm model 的路徑為 DATA/raw-data/\\[dataset\\]/spm\\[vocab_num\\].model"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "-q6qS0KhsqlO"
      },
      "source": [
        ""
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "N1F9jzmOsqlQ"
      },
      "source": [
        "### Binarize\n",
        "\n",
        "使用fairseq將資料轉為binary"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "u5miKBb5sqlQ"
      },
      "source": [
        "binpath = Path('./DATA/data-bin', mono_dataset_name)\n",
        "src_dict_file = './DATA/data-bin/ted2020/dict.en.txt'\n",
        "tgt_dict_file = src_dict_file\n",
        "monopref = str(mono_prefix/\"mono.tok\") # whatever filepath you get after applying subword tokenization\n",
        "if binpath.exists():\n",
        "    print(binpath, \"exists, will not overwrite!\")\n",
        "else:\n",
        "    !python -m fairseq_cli.preprocess\\\n",
        "        --source-lang 'zh'\\\n",
        "        --target-lang 'en'\\\n",
        "        --trainpref {monopref}\\\n",
        "        --destdir {binpath}\\\n",
        "        --srcdict {src_dict_file}\\\n",
        "        --tgtdict {tgt_dict_file}\\\n",
        "        --workers 2"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "p_xn2OKHsqlQ"
      },
      "source": [
        "### TODO: 生成反向翻譯資料\n",
        "\n",
        "將 binarized data 加入原本的資料夾中並用一個 split_name 取名\n",
        "\n",
        "ex. ./DATA/data-bin/ted2020/\\[split_name\\].zh-en.\\[\"en\", \"zh\"\\].\\[\"bin\", \"idx\"\\]\n",
        "\n",
        "便可以使用 generate_prediction(model, task, split=\"split_name\")來產生翻譯資料"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "PB5kEBCIsqlR"
      },
      "source": [
        "# 將 binarized data 加入原本的資料夾中並用一個 split_name 取名\n",
        "# ex. ./DATA/data-bin/ted2020/\\[split_name\\].zh-en.\\[\"en\", \"zh\"\\].\\[\"bin\", \"idx\"\\]\n",
        "!cp ./DATA/data-bin/mono/train.zh-en.zh.bin ./DATA/data-bin/ted2020/mono.zh-en.zh.bin\n",
        "!cp ./DATA/data-bin/mono/train.zh-en.zh.idx ./DATA/data-bin/ted2020/mono.zh-en.zh.idx\n",
        "!cp ./DATA/data-bin/mono/train.zh-en.en.bin ./DATA/data-bin/ted2020/mono.zh-en.en.bin\n",
        "!cp ./DATA/data-bin/mono/train.zh-en.en.idx ./DATA/data-bin/ted2020/mono.zh-en.en.idx"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "KqieDy5jsqlR"
      },
      "source": [
        "# hint: 用反向模型在 split='mono' 上進行預測，生成 prediction_file\n",
        "# generate_prediction( ... ,split=... ,outfile=... )"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "m7euUB3SsqlS"
      },
      "source": [
        "### TODO: 產生新的dataset\n",
        "\n",
        "1. 將翻譯出來的資料與原先的訓練資料結合\n",
        "2. 使用之前的spm model切出成Subword Units\n",
        "3. 重新使用fairseq將資料轉為binary"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "xC2bnOH0sqlS"
      },
      "source": [
        "# 合併剛剛生成的 prediction_file (.en) 以及中文 mono.zh (.zh)\n",
        "# \n",
        "# hint: 在此用剛剛的 spm model 對 prediction_file 進行切斷詞\n",
        "# spm_model.encode(line, out_type=str)\n",
        "# output: ./DATA/rawdata/mono/mono.tok.en & mono.tok.zh\n",
        "#\n",
        "# hint: 在此用 fairseq 把這些檔案再 binarize\n",
        "# binpath = Path('./DATA/data-bin/synthetic')\n",
        "# src_dict_file = './DATA/data-bin/ted2020/dict.en.txt'\n",
        "# tgt_dict_file = src_dict_file\n",
        "# monopref = ./DATA/rawdata/mono/mono.tok # or whatever path after applying subword tokenization, w/o the suffix (.zh/.en)\n",
        "# if binpath.exists():\n",
        "#     print(binpath, \"exists, will not overwrite!\")\n",
        "# else:\n",
        "#     !python -m fairseq_cli.preprocess\\\n",
        "#         --source-lang 'zh'\\\n",
        "#         --target-lang 'en'\\\n",
        "#         --trainpref {monopref}\\\n",
        "#         --destdir {binpath}\\\n",
        "#         --srcdict {src_dict_file}\\\n",
        "#         --tgtdict {tgt_dict_file}\\\n",
        "#         --workers 2"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "w9Zq_R8NsqlS"
      },
      "source": [
        "# 這裡用剛剛準備的檔案合併原先 ted2020 來生成最終 back-translation 的資料\n",
        "!cp -r ./DATA/data-bin/ted2020/ ./DATA/data-bin/ted2020_with_mono/\n",
        "\n",
        "!cp ./DATA/data-bin/synthetic/train.zh-en.zh.bin ./DATA/data-bin/ted2020_with_mono/train1.en-zh.zh.bin\n",
        "!cp ./DATA/data-bin/synthetic/train.zh-en.zh.idx ./DATA/data-bin/ted2020_with_mono/train1.en-zh.zh.idx\n",
        "!cp ./DATA/data-bin/synthetic/train.zh-en.en.bin ./DATA/data-bin/ted2020_with_mono/train1.en-zh.en.bin\n",
        "!cp ./DATA/data-bin/synthetic/train.zh-en.en.idx ./DATA/data-bin/ted2020_with_mono/train1.en-zh.en.idx"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "6BblDUe3sqlS"
      },
      "source": [
        "### TODO: 重新訓練\n",
        "\n",
        "當已經產生新的資料集\n",
        "\n",
        "1. 將實驗的參數設定表(config)中的datadir改為新的資料集(\"./DATA/data-bin/ted2020_with_mono\")\n",
        "2. 將實驗的參數設定表(config)中的source_lang與target_lang設定還原(\"en\", \"zh\")\n",
        "3. 將實驗的參數設定表(config)中的savedir更改(ex. \"./checkpoints/rnn-bt\")\n",
        "4. 重新訓練"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ERYLO2LJsqlT"
      },
      "source": [
        "# References"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "yTa0ndb7sqlT"
      },
      "source": [
        "1. <a name=ott2019fairseq></a>Ott, M., Edunov, S., Baevski, A., Fan, A., Gross, S., Ng, N., ... & Auli, M. (2019, June). fairseq: A Fast, Extensible Toolkit for Sequence Modeling. In Proceedings of the 2019 Conference of the North American Chapter of the Association for Computational Linguistics (Demonstrations) (pp. 48-53).\n",
        "2. <a name=vaswani2017></a>Vaswani, A., Shazeer, N., Parmar, N., Uszkoreit, J., Jones, L., Gomez, A. N., ... & Polosukhin, I. (2017, December). Attention is all you need. In Proceedings of the 31st International Conference on Neural Information Processing Systems (pp. 6000-6010).\n",
        "3. <a name=reimers-2020-multilingual-sentence-bert></a>Reimers, N., & Gurevych, I. (2020, November). Making Monolingual Sentence Embeddings Multilingual Using Knowledge Distillation. In Proceedings of the 2020 Conference on Empirical Methods in Natural Language Processing (EMNLP) (pp. 4512-4525).\n",
        "4. <a name=tiedemann2012parallel></a>Tiedemann, J. (2012, May). Parallel Data, Tools and Interfaces in OPUS. In Lrec (Vol. 2012, pp. 2214-2218).\n",
        "5. <a name=kudo-richardson-2018-sentencepiece></a>Kudo, T., & Richardson, J. (2018, November). SentencePiece: A simple and language independent subword tokenizer and detokenizer for Neural Text Processing. In Proceedings of the 2018 Conference on Empirical Methods in Natural Language Processing: System Demonstrations (pp. 66-71).\n",
        "6. <a name=sennrich-etal-2016-improving></a>Sennrich, R., Haddow, B., & Birch, A. (2016, August). Improving Neural Machine Translation Models with Monolingual Data. In Proceedings of the 54th Annual Meeting of the Association for Computational Linguistics (Volume 1: Long Papers) (pp. 86-96).\n",
        "7. <a name=edunov-etal-2018-understanding></a>Edunov, S., Ott, M., Auli, M., & Grangier, D. (2018). Understanding Back-Translation at Scale. In Proceedings of the 2018 Conference on Empirical Methods in Natural Language Processing (pp. 489-500).\n",
        "8. https://github.com/ajinkyakulkarni14/TED-Multilingual-Parallel-Corpus\n",
        "9. https://ithelp.ithome.com.tw/articles/10233122\n",
        "10. https://nlp.seas.harvard.edu/2018/04/03/attention.html"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "DiHzAjBPsqlT"
      },
      "source": [
        ""
      ],
      "execution_count": null,
      "outputs": []
    }
  ]
}